import numpy as np
import matplotlib.pyplot as plt
import itertools
import math
import scipy

def calculate_copeland(node1, nodes, edges):
    copeland_score = 0
    for node2 in nodes:
        if (node1, node2) in edges:
            copeland_score += 1
        # Assumes proper graph (i.e., if not u, v, then v, u)
    return copeland_score


def calculate_max_copeland(nodes, edges):
    return max(calculate_copeland(node, nodes, edges) for node in nodes)


def run_se(nodes, edges):
    """Returns Copeland score of winner on the seeding given by nodes
    Edges give indices into directions, and directions is a string consisting of 0s and 1s
    (0 indicates that the first item in the pair wins)
    """
    orig_nodes = nodes
    
    # Run SE
    while len(nodes) > 1:
        next_level = []
        for i in range(0, len(nodes), 2):
            if i + 1 == len(nodes):
                next_level.append(nodes[i])
            else:
                u, v = nodes[i], nodes[i + 1]
                if (u, v) in edges:
                    next_level.append(u)
                # Assumes proper graph (i.e., if not u, v, then v, u)
                else:
                    next_level.append(v)
        nodes = next_level
    
    return calculate_copeland(nodes[0], orig_nodes, edges)


# Uniform for now
# low and high are inclusive
def simulate_all_seedings(n, ):
    delta = (high - low) / (n - 1)
    scores = [i * delta for i in range(n)]
    nodes = list(range(n))
    rng = np.random.default_rng(seed)
    vals = rng.standard_normal(10)
    
    res = []
    for _ in range(ntrials):
        edges = set()
        for node1 in nodes:
            for node2 in nodes:
                if rng.normal(scores[node1], stddev) > rng.normal(scores[node2], stddev):
                    edges.add((node1, node2))
                else:
                    edges.add((node2, node1))
        total = 0
        n_seeds = 0
        for seeding in itertools.permutations(nodes):
            total += run_se(seeding, edges)
            n_seeds += 1
        res.append((total / n_seeds) / calculate_max_copeland(nodes, edges))
    return np.array(res)


# Uniform for now
# low and high are inclusive
def simulate(n, stddev, low, high, ntrials, nseedings):
    delta = (high - low) / (n - 1)
    scores = [i * delta for i in range(n)]
    nodes = list(range(n))
    rng = np.random.default_rng(0)
    vals = rng.standard_normal(10)
    
    res = []
    for _ in range(ntrials):
        edges = set()
        for node1 in nodes:
            for node2 in nodes:
                if rng.normal(scores[node1], stddev) > rng.normal(scores[node2], stddev):
                    edges.add((node1, node2))
                else:
                    edges.add((node2, node1))
        total = 0
        for _ in range(nseedings):
            total += run_se(rng.permutation(nodes), edges)
        res.append((total / nseedings) / calculate_max_copeland(nodes, edges))
    return np.array(res)

from time import time
def visualize(n, stddev, low, high, ntrials, nseedings, seed):
    start = time()
    simulated = simulate(n=n, stddev=stddev, low=low, high=high, ntrials=ntrials, nseedings=nseedings, seed=seed)
    print(time() - start, "to simulate")

    fig, ax = plt.subplots()
    ax.plot(simulated)
    ax.set_title(f"$n={n} \in [{low}, {high}]$, $\delta={stddev}$, ntrials={ntrials}, nseedings={nseedings}")
    plt.show()

    fig, ax = plt.subplots()
    ax.plot(np.sort(simulated))
    ax.set_title(f"(sorted) $n={n} \in [{low}, {high}]$, $\delta={stddev}$, ntrials={ntrials}, nseedings={nseedings}")
    plt.show()
    return simulated

simulated1 = visualize(n=8, stddev=1, low=0, high=10, ntrials=1000, nseedings=1000, seed=0)

28.79528784751892 to simulate

simulated2 = visualize(n=16, stddev=1, low=0, high=10, ntrials=1000, nseedings=1000, seed=0)

43.12546515464783 to simulate

simulated3 = visualize(n=32, stddev=1, low=0, high=10, ntrials=1000, nseedings=1000, seed=0)

72.55691504478455 to simulate

simulated4 = visualize(n=64, stddev=1, low=0, high=10, ntrials=1000, nseedings=1000, seed=0)

151.56541395187378 to simulate

simulated5 = visualize(n=128, stddev=1, low=0, high=10, ntrials=1000, nseedings=1000, seed=0)

410.7069180011749 to simulate

simulated6 = visualize(n=256, stddev=1, low=0, high=10, ntrials=1000, nseedings=1000, seed=0)

1152.8647711277008 to simulate

simulated_stddev_5 = []
for n in [8, 16, 32, 64, 128, 256]:
    simulated_stddev_5.append(visualize(n=n, stddev=5, low=0, high=10, ntrials=1000, nseedings=1000, seed=0))

28.935423851013184 to simulate

43.34131121635437 to simulate

76.04313898086548 to simulate